#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2025-08-18


import os
import threading
import time
import cv2
import uuid
import os
import platform
import argparse
from typing import Union
from queue import Queue
from spirems import Publisher, Subscriber, cvimg2sms, QoS, BaseNode, get_extra_args, Rate
import numpy as np


class LidarBinDirReaderNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        **kwargs
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=True,
            **kwargs
        )
        self.bin_dir = self.get_param("bin_dir", "/home/amov/spirecv-training-pro/mmdetection3d/demo/data/kitti")
        self.load_dim = self.get_param("load_dim", 4)
        self.use_dim = self.get_param("use_dim", [0, 1, 2, 3])
        self.fields = self.get_param("fields", ["x", "y", "z", "intensity"])
        self.frame_id = self.get_param("frame_id", "lidar_top")
        self.repeated = self.get_param("repeated", 0)
        self.auto_next = self.get_param("auto_next", False)
        self.rate = self.get_param("rate", 10)
        self.params_help()

        self.client_id = str(uuid.uuid4()).replace('-', '_')
        self.bin_i_queue = Queue()
        self.bin_i = 0
        self.done = False

        self._next_reader = Subscriber(
            '/' + job_name + '/launch_next', 'std_msgs::Boolean', self.launch_next,
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._bin_writer = Publisher(
            '/' + job_name + '/sensor/pointcloud', 'memory_msgs::PointCloud',
            ip=ip, port=port, qos=QoS.Reliability
        )
        
        files_and_directories = os.listdir(self.bin_dir)
        files_and_directories = sorted(files_and_directories)
        self.bin_fns = []
        for bin_name in files_and_directories:
            if os.path.splitext(bin_name)[-1] in ['.bin']:
                self.bin_fns.append(os.path.join(self.bin_dir, bin_name))
        self.start()

    def release(self):
        BaseNode.release(self)
        self._bin_writer.kill()
        self._next_reader.kill()

    def launch_next(self, msg: dict = None):
        if (isinstance(msg, dict) and msg['data']) or msg is None:
            if self.bin_i < len(self.bin_fns):
                self.bin_i_queue.put(self.bin_i)
                self.bin_i += 1
                if self.repeated and self.bin_i == len(self.bin_fns):
                    self.bin_i = 0
            else:
                self.done = True

    def run(self):
        while self.is_running():
            bin_i = self.bin_i_queue.get(block=True)
            if bin_i is None:
                break

            with open(self.bin_fns[bin_i], 'rb') as f:
                pts_bytes = f.read()
                points = np.frombuffer(pts_bytes, dtype=np.float32)

            points = points.reshape(-1, self.load_dim)
            points = points[:, self.use_dim]
            # print(points.shape)
            # print(points.dtype)

            msg = self._bin_writer.pcl2sms_mem(
                points, 
                fields=self.fields,     # fields为pcl每一列数据的具体描述
                frame_id=self.frame_id, # frame_id 可自定义设置
                timestamp=time.time()   # timestamp 可自定义设置
            )
            msg['pcl_id'] = bin_i
            msg['pcl_total'] = len(self.bin_fns)
            msg['file_name'] = os.path.basename(self.bin_fns[bin_i])
            msg['client_id'] = self.client_id
            self._bin_writer.publish(msg)

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = LidarBinDirReaderNode(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.launch_next()
    if node.auto_next:
        r = Rate(node.rate)
        while not node.done:
            node.launch_next()
            r.sleep()
        node.shutdown()
